#include <Windows.h>
#include <string>
#include <iostream>
#include <sstream>
#include <vector>
#include <format>
#include "Global_Include.h"
#include "Diskpart.h"
//#include "includes/SimpleDiskPartAPI.h"

//#error TODO：1.在进行卷压缩时，加入卷类型检测，使分区只能创建在FIXED上；无需加入对GPT分区的检测，在MBR上加载与在GPT商家在都一样

std::vector<WCHAR> GetUsedLetter();
std::vector<WCHAR> GetAvaiableLetterW();
DWORD GetDiskCount(
	_Out_ DWORD* dDiskNumber
);


/// <summary>
/// 创建PE分区
/// </summary>
/// <param name="SizeMB">分区大小（MB）</param>
/// <param name="PePartitionLetter">分配的PE分区盘符</param>
/// <param name="ShrinkPartition">被压缩分区盘符，若为压缩则为空</param>
/// <returns>成功返回true，失败返回false</returns>
bool CreatePeDiskW(
	_In_ unsigned long long* SizeMB,
	_Out_ WCHAR* PePartitionLetter,
	_Out_ WCHAR* ShrinkPartition
) {
	AstLog_WriteMessage(LogHandle, L"[RUNTIME][CREATEPE][DISKPART]开始创建PE分区");
	AstLog_WriteMessage(LogHandle, std::format(L"[RUNTIME][CREATEPE][DISKPART]创建分区大小：{} MB", *SizeMB));

	//获取可用盘符
	auto AvaiableLetter = GetAvaiableLetterW();
	//没有，报错
	if (AvaiableLetter.size() == 0) {
		AstThrowError_ThrowError(
			ERROR_RUNTIME_CREATEPE_DISKPART_NO_MORE_AVAIABLE_LETTER,
			L"[RUNTIME][CREATEPE][DISKPART]没有多余的盘符可供PE分区挂载"
		);
		return false;
	}
	AstLog_WriteMessage(LogHandle, std::format(L"[RUNTIME][CREATEPE][DISKPART]已获取到 {} 个可用盘符", AvaiableLetter.size()));

	//获取磁盘总数
	DWORD DiskNumber = 0;
	DWORD dStatus = GetDiskCount(
		&DiskNumber
	);
	//没有，报错
	if (dStatus != 0) {
		AstThrowError_ThrowError(
			ERROR_RUNTIME_CREATEPE_DISKPART_GET_HARDDISK_NUMBER_FAIL,
			std::format(L"[RUNTIME][CREATEPE][DISKPART]获取计算机中磁盘总数失败，WinErr代码：{}", dStatus)
		);
		return false;
	}
	AstLog_WriteMessage(
		LogHandle,
		std::format(L"[RUNTIME][CREATEPE][DISKPART]计算机中共有 {} 个硬盘", DiskNumber)
	);
	//处理盘符
	auto LetterW = AvaiableLetter[0];
	std::wstring LetterWstr{ LetterW };
	CHAR LetterA = _wstr2str(LetterWstr)[0];

	//开始创建分区
	//Step1 循环遍历分区
	AstLog_WriteMessage(
		LogHandle,
		L"[RUNTIME][CREATEPE][DISKPART]开始使用方案1创建分区！"
	);
	for (int i = 0; i < DiskNumber; i++) {
		AstLog_WriteMessage(
			LogHandle,
			std::format(L"[RUNTIME][CREATEPE][DISKPART]尝试在硬盘 {} 上创建分区", i)
		);

		bool bStatus = AstDiskPart_CreateVolume(i, LetterA, *SizeMB);
		if (bStatus) {
			AstLog_WriteMessage(
				LogHandle,
				std::format(L"[RUNTIME][CREATEPE][DISKPART]在硬盘 {} 上创建分区成功，盘符：{}，无压缩分区！", i, LetterA)
			);
			*PePartitionLetter = LetterW;
			return true;
		}
		AstLog_WriteMessage(
			LogHandle,
			std::format(L"[RUNTIME][CREATEPE][DISKPART]在硬盘 {} 上创建分区失败", i)
		);

	}

	//Step2  压缩分区创建分区
	AstLog_WriteMessage(
		LogHandle,
		std::format(L"[RUNTIME][CREATEPE][DISKPART]使用方案1创建PE分区失败，尝试使用方案2创建PE分区")
	);

	//获取使用盘符
	auto UsedLetter = GetUsedLetter();
	if (UsedLetter.size() == 0) {
		AstThrowError_ThrowError(
			ERROR_RUNTIME_CREATEPE_DISKPART_NO_USED_LETTER,
			L"[RUNTIME][CREATEPE][DISKPART]没有正在使用的盘符！"
		);
		return false;
	}

	for (auto Letter : UsedLetter) {

		//组装盘符
		std::wstring MountedLetterWstr{ Letter };
		CHAR MountedLetterA = _wstr2str(MountedLetterWstr)[0];

		AstLog_WriteMessage(
			LogHandle,
			std::format(L"[RUNTIME][CREATEPE][DISKPART]尝试压缩分区“{}”", MountedLetterWstr)
		);

		bool Status = AstDiskPart_ShrinkVolume(MountedLetterA, *SizeMB);

		if (!Status) {
			AstLog_WriteMessage(
				LogHandle,
				std::format(L"[RUNTIME][CREATEPE][DISKPART]压缩分区“{}”失败", MountedLetterWstr)
			);
		}

		for (int i = 0; i < DiskNumber; i++) {
			AstLog_WriteMessage(
				LogHandle,
				std::format(L"[RUNTIME][CREATEPE][DISKPART]尝试在硬盘 {} 上创建分区", i)
			);

			bool bStatus = AstDiskPart_CreateVolume(i, LetterA, *SizeMB);
			if (bStatus) {
				AstLog_WriteMessage(
					LogHandle,
					std::format(L"[RUNTIME][CREATEPE][DISKPART]在硬盘 {} 上创建分区成功，盘符：{}，压缩分区：{}", i, LetterA, MountedLetterWstr)
				);

				*PePartitionLetter = LetterW;
				*ShrinkPartition = Letter;

				return true;
			}
			AstLog_WriteMessage(
				LogHandle,
				std::format(L"[RUNTIME][CREATEPE][DISKPART]在硬盘 {} 上创建分区失败", i)
			);

		}
	}

	return false;
}

std::vector<WCHAR> GetUsedLetter() {
	std::vector<WCHAR> DriveLetter;

	//创建完全盘符列表
	std::vector<WCHAR> AllLetter;
	for (WCHAR wc = L'C'; wc <= L'Z'; wc++) {
		std::wstringstream LetterPath;
		LetterPath << L"\\\\.\\" << wc << L":";

		//尝试打开磁盘
		HANDLE hDevice = CreateFile(
			LetterPath.str().c_str(),
			GENERIC_READ,
			FILE_SHARE_READ | FILE_SHARE_WRITE,
			NULL,
			OPEN_EXISTING,
			0,
			NULL
		);

		if (hDevice != INVALID_HANDLE_VALUE) {
			CloseHandle(hDevice);
			DriveLetter.push_back(wc);
		}

	}

	if (DriveLetter.size() == 0) {
		return (std::vector<WCHAR>)NULL;
	}

	SetLastError(0);
	return DriveLetter;
}

std::vector<WCHAR> GetAvaiableLetterW() {
	std::vector<WCHAR> DriveLetter;

	//创建完全盘符列表
	std::vector<WCHAR> AllLetter;
	for (WCHAR wc = L'C'; wc <= L'Z'; wc++) {
		std::wstringstream LetterPath;
		LetterPath << L"\\\\.\\" << wc << L":";

		//尝试打开磁盘
		HANDLE hDevice = CreateFile(
			LetterPath.str().c_str(),
			GENERIC_READ,
			FILE_SHARE_READ | FILE_SHARE_WRITE,
			NULL,
			OPEN_EXISTING,
			0,
			NULL
		);

		if (hDevice == INVALID_HANDLE_VALUE) {
			CloseHandle(hDevice);
			DriveLetter.push_back(wc);
		}
		else if (GetLastError() == ERROR_FILE_NOT_FOUND) {
			return (std::vector<WCHAR>)NULL;
		}

	}

	if (DriveLetter.size() == 0) {
		return (std::vector<WCHAR>)NULL;
	}

	SetLastError(0);
	return DriveLetter;
}

DWORD GetDiskCount(
	_Out_ DWORD* dDiskNumber
) {
	*dDiskNumber = 0;
	while (true) {
		// 拼接硬盘路径，拼接后路径 \\\\.\\PhysicalDriveX (X是硬盘号)
		std::wstringstream wss_PhysicalDrivePath;
		wss_PhysicalDrivePath << L"\\\\.\\PhysicalDrive";
		wss_PhysicalDrivePath << *dDiskNumber;

		HANDLE hDevice = INVALID_HANDLE_VALUE; //硬盘句柄
		//打开磁盘
		hDevice = CreateFile(
			wss_PhysicalDrivePath.str().c_str(),
			GENERIC_READ,
			FILE_SHARE_READ | FILE_SHARE_WRITE,
			NULL,
			OPEN_EXISTING,
			0,
			NULL
		);


		if (hDevice == INVALID_HANDLE_VALUE) {
			DWORD ErrorCode = GetLastError();
			//没有磁盘，跳出循环
			if (ErrorCode == ERROR_FILE_NOT_FOUND) break;

			//读取失败，报错返回
			else {
				*dDiskNumber = -1;
				return ErrorCode;
			}
		}

		//否则，磁盘号递增
		CloseHandle(hDevice);
		*dDiskNumber = *dDiskNumber + 1;
	}

	return 0;
}